import asyncio
from typing import List, Dict, Any
from managers.vessel_manager import VesselManager
from models.schemas import AnalysisResponse
from managers.model_manager import LesionModelManager, MainVesselModelManager, SegmentationModelManager, SegmentNamedModelManager, F_SegmentNamedModelManager, LS_SegmentNamedModelManager, RCA_SegmentNamedModelManager, PositionEstimationModelManager
import tritonclient.http as httpclient
from core.config import TRITON_SERVER_URL
import SimpleITK as sitk
import time
import os
from loguru import logger
import numpy as np

async def get_ai_result(dicom_uid: str, img_data: List[List[str]], seg_data: List[List[str]], dsa_position: str) -> Dict[str, Any]:
    # 使用 asyncio.sleep 而不是 time.sleep，避免阻塞事件循环
    # 模拟AI处理时间，实际使用时应该替换为真实的AI推理代码
    # await asyncio.sleep(5)  # 改为5秒用于测试，实际可根据需要调整
    client = httpclient.InferenceServerClient(url=TRITON_SERVER_URL)
    img_data = np.array(img_data)
    seg_data = np.array(seg_data)
    img_data = (img_data - img_data.min()) / (img_data.max() - img_data.min()) * 255
    img_data = img_data.astype(np.uint8)
    seg_data = seg_data.astype(np.uint8)
    logger.info(f"img_data: {img_data.shape}, seg_data: {seg_data.shape}, dsa_position: {dsa_position}")
    local_date = time.strftime("%Y%m%d")
    visulize_path = 'visulization' +'/' + local_date + '/' + dicom_uid
    os.makedirs(visulize_path, exist_ok=True)
    start_time = time.time()
    # 1/2/3. 并发执行检测、命名模型和主支分割模型
    main_vessel_task = asyncio.create_task(run_main_vessel_pipeline(client, seg_data))
    lesion_task = asyncio.create_task(run_lesion_pipeline(client, img_data))
    named_task = asyncio.create_task(run_named_pipeline(client, img_data, dsa_position))
    if dsa_position == 'LS' or dsa_position == 'F':
        main_vessel_result, lesion_result, named_result = await asyncio.gather(main_vessel_task, lesion_task, named_task)
    else:
        lesion_result, named_result = await asyncio.gather(lesion_task, named_task)
        main_vessel_result = None
    logger.debug(f"named_result: {named_result.max()}")
    end_time = time.time()
    logger.info(f"并发执行检测与命名模型推理成功，耗时: {end_time - start_time}")
    
    # 4. 检测结果融合处理
    vessel_manager = VesselManager()
    
    # 4.1 设置vessel信息
    vessel_msg_dict = {
        'img_data': img_data, 
        'seg_data': seg_data, 
        'dsa_position': dsa_position, 
        'lesion_result': lesion_result, 
        'named_result': named_result,
        'main_vessel_result': main_vessel_result
    }
    await vessel_manager.set_vessel_message(vessel_msg_dict)
    
    # 4.2 获取skeleton, 并设置到vessel_manager中
    await vessel_manager.get_skeleton()
    # 中心线获取和可视化可以并发执行，因为它们都依赖已完成的骨架
    centerline_result = await vessel_manager.get_vessel_centerline()
    await vessel_manager.visulize_vessel_msg(centerline_result, visulize_path)
    end_time = time.time()
    logger.info(f"血管处理流程完成，总耗时: {end_time - start_time}")
    
    ai_result = await vessel_manager.return_result()

    return ai_result

async def run_lesion_pipeline(client, image_data): 
    lesion_manager = LesionModelManager()
    input_data = await lesion_manager.preprocess(image_data)
    result = await lesion_manager.run(input_data, client, model_version="1")
    return await lesion_manager.postprocess(result)

async def run_main_vessel_pipeline(client, image_data):
    main_vessel_manager = MainVesselModelManager()
    input_data = await main_vessel_manager.preprocess(image_data)
    result = await main_vessel_manager.run(input_data, client, model_version="1")
    return await main_vessel_manager.postprocess(result)

async def run_named_pipeline(client, image_data, position_label):
    if position_label == 'F':
        logger.info("执行分割模型：F左冠足位")
        named_manager = F_SegmentNamedModelManager()
    elif position_label == 'LS':
        logger.info("执行分割模型：LS左冠左肩位")
        named_manager = LS_SegmentNamedModelManager()
    elif position_label == 'RCA':
        logger.info("执行分割模型：RCA右冠体位")
        named_manager = RCA_SegmentNamedModelManager()
    else:
        raise ValueError(f"Invalid DSA position: {position_label}")
    input_data = await named_manager.preprocess(image_data)
    result = await named_manager.run(input_data, client, model_version="1")
    return await named_manager.postprocess(result)